#include "Algorithm.hpp"

void Algorithm::run_private(cv::Mat orig, cv::Mat mask, cv::Mat &result, int windowSize)
{
	// windowSize must be odd
	assert(windowSize % 2 == 1);

	int offset = windowSize / 2;
	cv::Mat image_padded, mask_padded, confidence_padded, data_paddeed;

	initMat(orig, mask, image_padded, mask_padded, confidence_padded, data_paddeed, result, offset);

	this->precedure(image_padded, mask_padded);

	int iter = 0;
	while (true)
	{
		cv::Mat dx, dy, nx, ny, fillFront_padded;

		update_gradients(image_padded, mask_padded, dx, dy);
		update_fillfront(mask_padded, fillFront_padded, nx, ny);

		//fillfront comes to the end
		if (cv::sum(fillFront_padded)[0] == 0)
		{
			if (cv::sum(mask)[0] > 0)
			{
				//remaining mask region is 1 pixel thick.
				//fillFront_padded = mask_padded;
				//FIXME: this is temporary. try to find a way around this.
			}
			break;
		}

		// calculate the point on the fill front with the maximum priority
		const cv::SparseMat::Node *p = calculateMaxPriorityPoint(fillFront_padded, confidence_padded, data_paddeed, dx, dy, nx, ny, windowSize, 255);
		int i = p->idx[0], j = p->idx[1];

		std::vector<std::pair<float, cv::Point2i>> candidators;

		this->findCandidatePatch(image_padded, mask_padded, i, j, windowSize, candidators); //cv::noArray()

		cv::Point2i best_ij = this->findBestPatch(image_padded, mask_padded, i, j, windowSize, candidators, max_edge_vector(dx, dy, i, j, offset));

		patch(image_padded, mask_padded, i, j, windowSize, best_ij);

		cv::Mat maskTemplate = mask_padded.rowRange(i - offset, i + offset + 1).colRange(j - offset, j + offset + 1);
		//update confidence
		confidence_padded.rowRange(i - offset, i + offset + 1).colRange(j - offset, j + offset + 1).setTo(confidence_padded.at<float>(i, j), maskTemplate);
		//updata mask
		mask_padded.rowRange(i - offset, i + offset + 1).colRange(j - offset, j + offset + 1).setTo(cv::Scalar(0));

		++iter;
		std::cout << iter << std::endl;

		this->afterIteration(image_padded, mask_padded, confidence_padded, data_paddeed, i, j);

		//visualize
		this->visualize(image_padded, mask_padded, windowSize);
		cv::waitKey(1);
	}
	//copy the non-padded parts of image_padded to result
	image_padded.rowRange(offset, image_padded.rows - offset).colRange(offset, image_padded.cols - offset).copyTo(result);
}

cv::Mat Algorithm::calculateFillFront(cv::Mat &mask)
{
	cv::Mat diskElem = cv::Mat::ones(cv::Size(3, 3), CV_8UC1);
	diskElem.at<unsigned char>(0, 0) = 0;
	diskElem.at<unsigned char>(0, 2) = 0;
	diskElem.at<unsigned char>(2, 0) = 0;
	diskElem.at<unsigned char>(2, 2) = 0;

	cv::Mat result, maskEroded;

	cv::erode(mask, maskEroded, diskElem);
	cv::subtract(mask, maskEroded, result);
	return result;
}

cv::SparseMat Algorithm::calculateSparseFillFront(cv::Mat &fillFront)
{
	cv::SparseMat result(fillFront);
	return result;
}

const cv::SparseMat::Node *Algorithm::calculateMaxPriorityPoint(cv::Mat &fillFront_padded, cv::Mat &confidence_padded, cv::Mat &data_paddeed, cv::Mat &dx, cv::Mat &dy, cv::Mat &nx, cv::Mat &ny, int windowSize, float alpha)
{
	float max_priority = -1;
	const cv::SparseMat::Node *max_point;

	int offset = windowSize / 2;
	cv::SparseMat sparseFillFront = calculateSparseFillFront(fillFront_padded);
	cv::SparseMatConstIterator
		it = sparseFillFront.begin(),
		it_end = sparseFillFront.end();

	//lets set the results to the first point first
	//later if we find better points it will change
	max_point = it.node();

	for (; it != it_end; ++it)
	{
		const cv::SparseMat::Node *node = it.node();
		int i = node->idx[0];
		int j = node->idx[1];

		//calculate data terms
		float data_term = calculateDataTerm(dx, dy, nx, ny, i, j, offset, alpha);
		data_paddeed.at<float>(i, j) = data_term;

		//calculate confidence terms
		float conf_term = calculateConfidenceTerm(confidence_padded, i, j, windowSize, offset);
		confidence_padded.at<float>(i, j) = conf_term;

		float priority = this->calculatePriority(data_term, conf_term);

		if (priority > max_priority)
		{
			max_priority = priority;
			max_point = node;
		}
	}

	return max_point;
}

float Algorithm::calculateDataTerm(cv::Mat &dx, cv::Mat &dy, cv::Mat &nx, cv::Mat &ny, int i, int j, int offset, float alpha)
{
	double data_term;

	cv::Point2f data_ij = max_edge_vector(dx, dy, i, j, offset);

	double abs_n = abs_vector(nx.at<float>(i, j), ny.at<float>(i, j));
	data_term = abs(data_ij.x * nx.at<float>(i, j) + data_ij.y * ny.at<float>(i, j));
	
	data_term /= abs_n;
	data_term /= alpha;

	return data_term;
}

float Algorithm::calculateConfidenceTerm(cv::Mat &confidence, int i, int j, int windowSize, int offset)
{
	float conf_term = 0;

	for (int k = -offset; k <= offset; ++k)
	{
		for (int l = -offset; l <= offset; ++l)
		{
			conf_term += confidence.at<float>(i + k, j + l);
		}
	}
	conf_term /= (float)(windowSize * windowSize);
	return conf_term;
}

void Algorithm::visConfidence(cv::Mat &confidence, cv::Mat &orig, cv::Mat &mask, int windowSize)
{
	cv::Mat origFloat, image_padded, mask_padded, mask_max_1;
	orig.convertTo(origFloat, CV_32FC3);
	int offset = windowSize / 2;
	cv::copyMakeBorder(origFloat, image_padded, offset, offset, offset, offset, cv::BORDER_CONSTANT, cv::Scalar(0));
	cv::copyMakeBorder(mask, mask_padded, offset, offset, offset, offset, cv::BORDER_CONSTANT, cv::Scalar(0));
	mask_padded.convertTo(mask_max_1, CV_8UC1, 1. / 255, 0);
	cv::subtract(cv::Scalar(1), mask_max_1, confidence);
	confidence.convertTo(confidence, CV_32FC1);

	for (int x = 0; x != mask_padded.rows; ++x)
	{
		for (int y = 0; y != mask_padded.cols; ++y)
		{
			if (mask_padded.at<unsigned char>(x, y) != 0)
			{
				float conf_term = this->calculateConfidenceTerm(confidence, x, y, windowSize, offset);
				confidence.at<float>(x, y) = conf_term;
			}
		}
	}
}

void Algorithm::visDataTerm(cv::Mat &dataTerm, cv::Mat &orig, cv::Mat &mask, int windowSize, float alpha)
{
	cv::Mat origFloat, maskInv, fillFront, image_padded, mask_padded, mask_max_1, confidence;
	orig.convertTo(origFloat, CV_32FC3);
	int offset = windowSize / 2;
	cv::copyMakeBorder(origFloat, image_padded, offset, offset, offset, offset, cv::BORDER_CONSTANT, cv::Scalar(0));
	cv::copyMakeBorder(mask, mask_padded, offset, offset, offset, offset, cv::BORDER_CONSTANT, cv::Scalar(0));

	cv::Mat dx, dy, nx, ny, image_padded_gray;
	cv::cvtColor(image_padded, image_padded_gray, cv::COLOR_BGR2GRAY);
	cv::Sobel(image_padded_gray, dx, CV_32F, 1, 0, 3, 1, 0, cv::BORDER_CONSTANT);
	cv::Sobel(image_padded_gray, dy, CV_32F, 0, 1, 3, -1, 0, cv::BORDER_CONSTANT);
	cv::Sobel(mask_padded, nx, CV_32F, 1, 0, 3, 1, 0, cv::BORDER_CONSTANT);
	cv::Sobel(mask_padded, ny, CV_32F, 0, 1, 3, 1, 0, cv::BORDER_CONSTANT);

	dataTerm = cv::Mat::zeros(image_padded.rows, image_padded.cols, CV_32FC1);

	for (int x = 0; x != mask_padded.rows; ++x)
	{
		for (int y = 0; y != mask_padded.cols; ++y)
		{
			if (mask_padded.at<unsigned char>(x, y) != 0)
			{
				float d_term = this->calculateDataTerm(dx, dy, nx, ny, x, y, offset, alpha);
				dataTerm.at<float>(x, y) = d_term;
			}
		}
	}
}

void Algorithm::update_gradients(cv::Mat &image_padded, cv::Mat &mask_padded, cv::Mat &dx, cv::Mat &dy)
{
	cv::Mat image_padded_gray;

	// convert image_padded to grayscale and calculate gradients
	//note that we attemp to get (dy,-dx)
	cv::cvtColor(image_padded, image_padded_gray, cv::COLOR_BGR2GRAY);
	cv::Sobel(image_padded_gray, dx, CV_32F, 0, 1, 3, 1, 0, cv::BORDER_CONSTANT);
	cv::Sobel(image_padded_gray, dy, CV_32F, 1, 0, 3, -1, 0, cv::BORDER_CONSTANT);
	// set the dx and dy to zero in the masked area.
	dx.setTo(cv::Scalar(0), mask_padded);
	dy.setTo(cv::Scalar(0), mask_padded);
}

void Algorithm::update_fillfront(cv::Mat &mask_padded, cv::Mat &fillFront_padded, cv::Mat &nx, cv::Mat &ny)
{
	//calculate fillfront and (nx,ny)
	cv::Mat maskEroded;
	cv::Mat diskElem = cv::Mat::ones(cv::Size(3, 3), CV_8UC1);

	diskElem.at<unsigned char>(0, 0) = 0;
	diskElem.at<unsigned char>(0, 2) = 0;
	diskElem.at<unsigned char>(2, 0) = 0;
	diskElem.at<unsigned char>(2, 2) = 0;

	cv::erode(mask_padded, maskEroded, diskElem);
	cv::subtract(mask_padded, maskEroded, fillFront_padded);

	cv::Sobel(mask_padded, nx, CV_32F, 1, 0, 3, 1, 0, cv::BORDER_CONSTANT);
	cv::Sobel(mask_padded, ny, CV_32F, 0, 1, 3, 1, 0, cv::BORDER_CONSTANT);
}

void Algorithm::initMat(cv::Mat &orig, cv::Mat &mask, cv::Mat &image_padded, cv::Mat &mask_padded, cv::Mat &confidence_padded, cv::Mat &data_paddeed, cv::Mat &result, int offset)
{
	cv::Mat maskInv, fillFront_padded, confidence, data;

	mask.convertTo(maskInv, CV_8UC1, -1, 1);
	//initial confidence, confidence is (1 - mask) which is maskInv.
	maskInv.convertTo(confidence, CV_32FC1);
	//initial data,assign value as you like
	maskInv.convertTo(data, CV_32FC1);

	//clear the masked parts of the image,cv::Scalar(0, 0, 0)
	orig.setTo(cv::Scalar(255, 0, 255), mask);
	result = cv::Mat::zeros(orig.rows, orig.cols, CV_8UC3);

	//padding
	cv::copyMakeBorder(orig, image_padded, offset, offset, offset, offset, cv::BORDER_CONSTANT, cv::Scalar(0));
	cv::copyMakeBorder(mask, mask_padded, offset, offset, offset, offset, cv::BORDER_CONSTANT, cv::Scalar(0));
	cv::copyMakeBorder(confidence, confidence_padded, offset, offset, offset, offset, cv::BORDER_CONSTANT, cv::Scalar(1));
	cv::copyMakeBorder(data, data_paddeed, offset, offset, offset, offset, cv::BORDER_CONSTANT, cv::Scalar(1));
}

void Algorithm::patch(cv::Mat &image_padded, cv::Mat &mask_padded, int i, int j, int windowSize, cv::Point2i best_ij)
{
	int offset = windowSize / 2;

	cv::Range sourceRowRange_padded(best_ij.x - offset, best_ij.x + offset + 1);
	cv::Range sourceColRange_padded(best_ij.y - offset, best_ij.y + offset + 1);
	cv::Range targetRowRange_padded(i - offset, i + offset + 1);
	cv::Range targetColRange_padded(j - offset, j + offset + 1);

	cv::Mat maskTemplate = mask_padded.rowRange(targetRowRange_padded).colRange(targetColRange_padded);

	image_padded.rowRange(sourceRowRange_padded).colRange(sourceColRange_padded).copyTo(image_padded.rowRange(targetRowRange_padded).colRange(targetColRange_padded), maskTemplate);
}

//pHash算法
cv::Mat Algorithm::pHashValue(cv::Mat &image)
{
	cv::Mat floatImage, imageDct;
	image.convertTo(floatImage, CV_32FC1);
	cv::dct(floatImage, imageDct);
	cv::Scalar imageMean = mean(floatImage);
	return (floatImage > imageMean[0]);
}

int Algorithm::calHammingDistance(cv::Mat modelHashCode, cv::Mat testHashCode)
{
	if (modelHashCode.size() != testHashCode.size())
		return -1;
	return cv::countNonZero(modelHashCode != testHashCode);
}